
from translate.xiaomi_translate import XiaomiTranslator
from googletrans import Translator
from typing import List, Dict, Union
from pathlib import Path
import edge_tts
import tempfile
from utils.hash import sha1_hash_file
import json
from loguru import logger
import shutil
import os
from audios.audio_asr import AudioAsr


def extract_texts(video_or_audio_file: Union[Path, str], zh_lang: str='老挝语') -> List[Dict]:

    parser = AudioAsr()
    return parser.parse(video_or_audio_file, zh_lang)


def translate_texts(audio_texts: List[Dict], to_lan='Chinese') -> List[Dict]:
    """调用小米的翻译模型, 批量翻译语音的转录文本"""
    from nlp.lang_detect import LanguageDetector
    lang_detector = LanguageDetector()
    translator = XiaomiTranslator()

    results = []
    for item in audio_texts:
        text = item['text']
        lang = lang_detector.detect(text)
        translated = translator.translate(text, from_lan=lang, to_lan=to_lan)
        logger.info(translated)
        item['lang'] = lang
        item['translated'] = translated
        results.append(item)
    return results


async def google_translate_texts(audio_texts: List[Dict], to_lan='Chinese') -> List[Dict]:
    """调用google的翻译模型, 批量翻译语音的转录文本"""
    translator = Translator()

    results = []
    for item in audio_texts:
        text = item['text']
        # lang = lang_detector.detect(text)
        translated = await translator.translate(text, dest=to_lan)
        logger.info(translated.text)
        item['lang'] = "Lao"
        item['translated'] = translated.text
        results.append(item)
    return results


def tts(audio_texts: List[Dict], save_dir: Union[Path, str] = None) -> List[Dict]:
    """调用微软edge_tts合成语音
    Name: zh-CN-XiaoxiaoNeural
    Name: zh-CN-XiaoyiNeural
    Name: zh-CN-YunjianNeural
    Name: zh-CN-YunxiNeural
    Name: zh-CN-YunxiaNeural
    Name: zh-CN-YunyangNeural
    Name: zh-CN-liaoning-XiaobeiNeural
    Name: zh-CN-shaanxi-XiaoniNeural
    """
    save_dir = save_dir or Path(__file__).parent
    save_dir = Path(save_dir)
    save_dir.mkdir(exist_ok=True)

    results = []
    for item in audio_texts:
        communicate = edge_tts.Communicate(text=item['translated'],
            voice="zh-CN-XiaoyiNeural",
            rate='+0%',
            volume= '+0%',
            pitch= '+0Hz')
        fp, clip_file = tempfile.mkstemp(suffix='.wav')
        try:
            communicate.save_sync(clip_file)
        except Exception as e:
            logger.error(e)
            os.close(fp)
            continue
        os.close(fp)
        sha1_hash = sha1_hash_file(clip_file)
        target_file = (save_dir / f'{sha1_hash}.wav').as_posix()
        shutil.move(clip_file, target_file)
        item['tts'] = target_file
        results.append(item)
        logger.info(target_file)
    return results


if __name__ == '__main__':
    import asyncio
    with open('../audios/parse_result.json', 'r', encoding='utf-8') as fs:
        audio_texts = json.load(fs)
    print(len(audio_texts))
    logger.info('google翻译, start')
    results = asyncio.run(google_translate_texts(audio_texts,to_lan='zh'))
    logger.info('google翻译, end')
    with open('./translated_result.json', 'w', encoding='utf-8', newline='') as fw:
        json.dump(results, fw, ensure_ascii=False, indent=4)
    results = tts(results, save_dir='./tts')
    logger.info('tts, start')